# 这是一个示例 Python 脚本。

# 按 Shift+F10 执行或将其替换为您的代码。
# 按 双击 Shift 在所有地方搜索类、文件、工具窗口、操作和设置。

# from transformers import BertModel, BertConfig, BertTokenizer
import torch
from data_util.data_generate import Dataset
from data_util.data_pro import get_data

from train import train
from torch import nn

from model.longformer_model import LongForm
import model.config as conf
from torch.utils.data import DataLoader

import matplotlib.pyplot as plt

import os
# bert_model = BertModel.from_pretrained('bert_base_chinese')
# config = BertConfig.from_pretrained('bert_base_chinese')
# tokenizer = BertTokenizer.from_pretrained('bert_base_chinese')

#
# 按间距中的绿色按钮以运行脚本。
if __name__ == '__main__':
    # print_hi('PyCharm')
    torch.manual_seed(123)
    train_list = get_data(conf.train_path)
    valid_list = get_data(conf.valid_path)
    train_dataset = Dataset(train_list)
    valid_dataset = Dataset(valid_list)
    train_dataloader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=True)
    valid_dataloader = DataLoader(valid_dataset, batch_size=conf.batch_size, shuffle=True)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    # model = Bert_Attention()
    lf_model = LongForm(n_model=conf.n_model, cls_nu=2)
    lf_model.to(device)
    # bigbird_model = LongForm(n_model=conf.n_model, cls_nu=len(conf.cls))
    # bigbird_model.to(device)
    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.Adam(lf_model.parameters(), lr=conf.lr)
    # lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.95)
    if not os.path.exists(conf.checkpoint):
        os.makedirs(conf.checkpoint)
    # train(bigbird_model, train_dataloader, valid_dataloader, device, criterion, lr_scheduler, optimizer)
    train_loss_list, valid_loss_list, acc_pred_list = train(lf_model, train_dataloader, valid_dataloader, device, criterion, optimizer)
    x_list = range(conf.epochs)
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.title('生产环境训练结果')
    plt.plot(x_list, train_loss_list, label='train loss')
    plt.plot(x_list, valid_loss_list, label='valid loss')
    plt.plot(x_list, acc_pred_list, label='acc')
    plt.legend(loc='best')
    plt.savefig('image/result.png')

    pass
# 访问 https://www.jetbrains.com/help/pycharm/ 获取 PyCharm 帮助
